package com.small.nacos.naming.push;

import com.alibaba.nacos.common.utils.JacksonUtils;
import com.small.nacos.naming.consistency.KeyBuilder;
import com.small.nacos.naming.core.Service;
import com.small.nacos.naming.misc.GlobalExecutor;
import com.small.nacos.naming.misc.SwitchDomain;
import org.apache.commons.collections4.MapUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.ApplicationListener;
import org.springframework.stereotype.Component;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetSocketAddress;
import java.net.SocketException;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.zip.GZIPOutputStream;

/**
 * @Author zhoujin
 * @Date 2022/11/2 10:38
 */
@Component
public class PushService implements ApplicationContextAware, ApplicationListener<ServiceChangeEvent> {

    private static final Logger LOGGER = LoggerFactory.getLogger(PushService.class);

    @Autowired
    private SwitchDomain switchDomain;

    private ApplicationContext applicationContext;

    private static DatagramSocket udpSocket;

    private static ConcurrentMap<String, Future> futureMap = new ConcurrentHashMap<>();

    private static ConcurrentMap<String, ConcurrentHashMap<String, PushClient>> clientMap = new ConcurrentHashMap<>();

    private static ConcurrentMap<String, Receiver.AckEntry> ackMap = new ConcurrentHashMap<>();

    private static final int MAX_RETRY_TIMES = 1;

    private static final long ACK_TIMEOUT_NANOS = TimeUnit.SECONDS.toNanos(10L);

    static {
        try {
            udpSocket = new DatagramSocket();
            Thread thread = new Thread(new Receiver());
            thread.setDaemon(true);
            thread.setName(PushService.class.getName());
            thread.start();

        } catch (SocketException e) {
            LOGGER.error("[NACOS-PUSH] failed to init push service");
        }
    }

    @Override
    public void onApplicationEvent(ServiceChangeEvent event) {

        Service service = event.getService();

        String key = KeyBuilder.buildServiceKey(service.getNamespaceId(), service.getName());

        ScheduledFuture<?> future = GlobalExecutor.scheduleUdpSender(() -> {

            try {
                ConcurrentHashMap<String, PushClient> clients = clientMap.get(key);
                if (MapUtils.isEmpty(clients)) {
                    return;
                }
                long lastRefTime = System.nanoTime();
                for (PushClient client : clients.values()) {
                    if (client.zombie()) {
                        clients.remove(client.toString());
                        continue;
                    }
                    Receiver.AckEntry ackEntry = prepareAckEntry(client, lastRefTime);
                    udpPush(ackEntry);
                }


            } catch (Exception e) {
                LOGGER.error("[NACOS-PUSH] failed to push serviceName: {} to client, error: {}", service.getName(), e);
            } finally {
                futureMap.remove(key);
            }


        }, 1000, TimeUnit.MILLISECONDS);

        futureMap.put(key, future);
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    private static void udpPush(Receiver.AckEntry ackEntry) {

        if (ackEntry == null) {
            LOGGER.error("[NACOS-PUSH] ackEntry is null.");
            return;
        }

        //超过了最大重试次数
        if (ackEntry.getRetryTimes() > MAX_RETRY_TIMES) {
            ackMap.remove(ackEntry.key);
            return;
        }

        try {
            ackMap.put(ackEntry.key, ackEntry);
            udpSocket.send(ackEntry.origin);
            //这里过10秒后会在发送一次，正常情况下接收到后ackMap会被删除，就不会发送
            GlobalExecutor.scheduleRetransmitter(new Retransmitter(ackEntry),
                    TimeUnit.NANOSECONDS.toMillis(ACK_TIMEOUT_NANOS), TimeUnit.MILLISECONDS);
            return;
        } catch (IOException e) {
            LOGGER.error("[NACOS-PUSH] failed to push data: {} to client: {}, error: {}", ackEntry.data,
                    ackEntry.origin.getAddress().getHostAddress(), e);
        } finally {
            ackEntry.increaseRetryTime();
        }
    }


    private Receiver.AckEntry prepareAckEntry(PushClient client, long lastRefTime) throws Exception {
        Map<String, Object> data = new HashMap<String, Object>(3);
        try {
            data.put("type", "dom");
            data.put("data", client.getDataSource().getData(client));
            data.put("lastRefTime", lastRefTime);

            String key = getAckKey(client.getSocketAddr().getAddress().getHostAddress(), client.getSocketAddr().getPort(),
                    lastRefTime);

            String dataStr = JacksonUtils.toJson(data);
            byte[] dataBytes = dataStr.getBytes(StandardCharsets.UTF_8);
            dataBytes = compressIfNecessary(dataBytes);
            DatagramPacket packet = new DatagramPacket(dataBytes, dataBytes.length, client.socketAddr);

            Receiver.AckEntry ackEntry = new Receiver.AckEntry(key, packet);
            ackEntry.data = data;

            return ackEntry;
        } catch (Exception e) {
            LOGGER.error("[NACOS-PUSH] failed to prepare data: {} to client: {}, error: {}", data,
                    client.getSocketAddr(), e);
        }
        return null;
    }

    private static String getAckKey(String host, int port, long lastRefTime) {
        return host + "," + port + "," + lastRefTime;
    }

    private static byte[] compressIfNecessary(byte[] dataBytes) throws IOException {
        // enable compression when data is larger than 1KB
        int maxDataSizeUncompress = 1024;
        if (dataBytes.length < maxDataSizeUncompress) {
            return dataBytes;
        }

        ByteArrayOutputStream out = new ByteArrayOutputStream();
        GZIPOutputStream gzip = new GZIPOutputStream(out);
        gzip.write(dataBytes);
        gzip.close();

        return out.toByteArray();
    }


    public void serviceChanged(Service service) {
        // merge some change events to reduce the push frequency:
        if (futureMap
                .containsKey(KeyBuilder.buildServiceKey(service.getNamespaceId(), service.getName()))) {
            return;
        }

        this.applicationContext.publishEvent(new ServiceChangeEvent(this, service));
    }

    public void addClient(String namespaceId, String serviceName, String clusters, String agent, InetSocketAddress socketAddr, DataSource dataSource) {

        PushClient client = new PushClient(namespaceId, serviceName, clusters, agent, socketAddr, dataSource);
        addClient(client);
    }

    public void addClient(PushClient client) {
        // client is stored by key 'serviceName' because notify event is driven by serviceName change

        String serviceKey = KeyBuilder.buildServiceKey(client.getNamespaceId(), client.getServiceName());
        ConcurrentMap<String, PushClient> clients = clientMap.get(serviceKey);
        if (clients == null) {
            clientMap.putIfAbsent(serviceKey, new ConcurrentHashMap<>(1024));
            clients = clientMap.get(serviceKey);
        }

        PushClient oldClient = clients.get(client.toString());
        if (oldClient != null) {
            oldClient.refresh();
        } else {
            PushClient res = clients.putIfAbsent(client.toString(), client);
            if (res != null) {
                LOGGER.warn("client: {} already associated with key {}", res.getAddrStr(), res.toString());
            }
            LOGGER.debug("client: {} added for serviceName: {}", client.getAddrStr(), client.getServiceName());
        }
    }


    public class PushClient {

        private String namespaceId;

        private String serviceName;

        private String clusters;

        private String agent;

        private InetSocketAddress socketAddr;

        private DataSource dataSource;

        public long lastRefTime = System.currentTimeMillis();

        public PushClient(String namespaceId, String serviceName, String clusters, String agent,
                          InetSocketAddress socketAddr, DataSource dataSource) {
            this.namespaceId = namespaceId;
            this.serviceName = serviceName;
            this.clusters = clusters;
            this.agent = agent;
            this.socketAddr = socketAddr;
            this.dataSource = dataSource;
        }

        public DataSource getDataSource() {
            return dataSource;
        }

        public boolean zombie() {
            return System.currentTimeMillis() - lastRefTime > switchDomain.getPushCacheMillis();
        }

        @Override
        public String toString() {
            StringBuilder sb = new StringBuilder();
            sb.append("serviceName: ").append(serviceName).append(", clusters: ").append(clusters).append(", address: ")
                    .append(socketAddr).append(", agent: ").append(agent);
            return sb.toString();
        }

        public String getAgent() {
            return agent;
        }

        public String getAddrStr() {
            return socketAddr.getAddress().getHostAddress() + ":" + socketAddr.getPort();
        }

        public String getIp() {
            return socketAddr.getAddress().getHostAddress();
        }

        @Override
        public int hashCode() {
            return Objects.hash(serviceName, clusters, socketAddr);
        }

        @Override
        public boolean equals(Object obj) {
            if (!(obj instanceof PushClient)) {
                return false;
            }

            PushClient other = (PushClient) obj;

            return serviceName.equals(other.serviceName) && clusters.equals(other.clusters) && socketAddr
                    .equals(other.socketAddr);
        }

        public String getClusters() {
            return clusters;
        }

        public void setClusters(String clusters) {
            this.clusters = clusters;
        }

        public String getNamespaceId() {
            return namespaceId;
        }

        public void setNamespaceId(String namespaceId) {
            this.namespaceId = namespaceId;
        }

        public String getServiceName() {
            return serviceName;
        }

        public void setServiceName(String serviceName) {
            this.serviceName = serviceName;
        }

        public InetSocketAddress getSocketAddr() {
            return socketAddr;
        }

        public void refresh() {
            lastRefTime = System.currentTimeMillis();
        }


    }


    public static class Retransmitter implements Runnable {

        Receiver.AckEntry ackEntry;

        public Retransmitter(Receiver.AckEntry ackEntry) {
            this.ackEntry = ackEntry;
        }

        @Override
        public void run() {
            if (ackMap.containsKey(ackEntry.key)) {
                LOGGER.info("retry to push data, key: " + ackEntry.key);
                udpPush(ackEntry);
            }
        }
    }

    public static class Receiver implements Runnable {

        @Override
        public void run() {
            while (true) {
                byte[] buffer = new byte[1024 * 64];
                DatagramPacket packet = new DatagramPacket(buffer, buffer.length);
                try {
                    udpSocket.receive(packet);
                    String json = new String(packet.getData(), 0, packet.getLength(), StandardCharsets.UTF_8).trim();
                    AckPacket ackPacket = JacksonUtils.toObj(json, AckPacket.class);
                    InetSocketAddress socketAddress = (InetSocketAddress) packet.getSocketAddress();
                    String ip = socketAddress.getAddress().getHostAddress();
                    int port = socketAddress.getPort();
                    if (System.nanoTime() - ackPacket.lastRefTime > ACK_TIMEOUT_NANOS) {
                        LOGGER.warn("ack takes too long from {} ack json: {}", packet.getSocketAddress(), json);
                    }
                    String ackKey = getAckKey(ip, port, ackPacket.lastRefTime);
                    //确认发送成功，移除任务
                    AckEntry ackEntry = ackMap.remove(ackKey);
                    if (ackEntry == null) {
                        throw new IllegalStateException(
                                "unable to find ackEntry for key: " + ackKey + ", ack json: " + json);
                    }

                } catch (Throwable e) {
                    LOGGER.error("[NACOS-PUSH] error while receiving ack data", e);
                }
            }
        }


        public static class AckEntry {

            public AckEntry(String key, DatagramPacket packet) {
                this.key = key;
                this.origin = packet;
            }

            public void increaseRetryTime() {
                retryTimes.incrementAndGet();
            }

            public int getRetryTimes() {
                return retryTimes.get();
            }

            public String key;

            public DatagramPacket origin;

            private AtomicInteger retryTimes = new AtomicInteger(0);

            public Map<String, Object> data;
        }

        public static class AckPacket {

            public String type;

            public long lastRefTime;

            public String data;
        }
    }

}
